import os
import sys

from PIL import Image
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget, QPushButton, QHBoxLayout, \
    QMessageBox, QFileDialog
from PyQt5.QtGui import QImage, QPixmap
import cv2
from ultralytics import YOLO
from train import recognition


class Worker:
    def __init__(self):
        self.model = None

    def load_model(self):
        model_path, _ = QFileDialog.getOpenFileName(None, "选择模型文件", "", "模型文件 (*.pt)")
        if model_path:
            self.model = YOLO(model_path)
            if self.model:
                return True
            else:
                return False

    def detect_image(self, image_path):
        image = cv2.imread(image_path)
        if image is not None:
            results = self.model.predict(image)
            return results
        else:
            return None


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Numb753")
        self.setGeometry(100, 100, 800, 600)

        self.label = QLabel()
        self.label.setAlignment(Qt.AlignCenter)

        layout = QVBoxLayout()
        layout.addWidget(self.label)

        central_widget = QWidget()
        central_widget.setLayout(layout)
        self.setCentralWidget(central_widget)

        self.worker = Worker()

        # 创建按钮布局
        hbox_buttons = QHBoxLayout()
        layout.addLayout(hbox_buttons)

        # 添加模型选择按钮
        self.load_model_button = QPushButton("模型选择")
        self.load_model_button.clicked.connect(self.load_model)
        hbox_buttons.addWidget(self.load_model_button)

        # 添加图片检测按钮
        self.image_detect_button = QPushButton("图片检测")
        self.image_detect_button.clicked.connect(self.detect_image)
        self.image_detect_button.setEnabled(False)
        hbox_buttons.addWidget(self.image_detect_button)

        # 添加退出按钮
        self.exit_button = QPushButton("退出")
        self.exit_button.clicked.connect(self.exit_application)
        hbox_buttons.addWidget(self.exit_button)

    def detect_image(self):
        image_path, _ = QFileDialog.getOpenFileName(None, "选择图片文件", "", "图片文件 (*.jpg *.jpeg *.png)")
        if image_path:
            results = self.worker.detect_image(image_path)
            if results:
                # print(image_path)
                annotated_image = results[0].plot()
                height, width, channel = annotated_image.shape
                bytesPerLine = 3 * width
                qimage = QImage(annotated_image.data, width, height, bytesPerLine, QImage.Format_BGR888)
                pixmap = QPixmap.fromImage(qimage)
                self.label.setPixmap(pixmap.scaled(self.label.size(), Qt.KeepAspectRatio))
                boxes = results[0].boxes.cpu().numpy()
                xyxys = boxes.xyxy
                index = 0
                message = []
                for xyxy in xyxys:
                    x1, y1, x2, y2 = xyxy
                    box1 = (x1, y1, x2, y2)
                    img = Image.open(image_path)
                    img = img.crop(box1)
                    root = "D:/PycharmProjects/yolo_v8/cache"
                    path = root + "/" + "{}.jpg".format(index)
                    index = index + 1
                    img.save(path)
                    result = recognition(path)
                    print(result)
                    os.remove(path)
                    message.append(result)
                combined_message = '\n'.join(' '.join(item) for item in message)
                message_box = QMessageBox()
                message_box.setWindowTitle('{}'.format(image_path))
                message_box.setText('{}'.format(combined_message))
                # 设置消息框的图标（这里使用信息图标）
                message_box.setIcon(QMessageBox.Information)
                message_box.exec_()
        else:
            QMessageBox.critical(self, "错误", "请选择图片文件")

    def load_model(self):
        if self.worker.load_model():
            self.image_detect_button.setEnabled(True)

    def exit_application(self):
        # 终止程序运行
        sys.exit()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
